package com.khalil.kkadmin.module.ai.agent;

import com.khalil.kkadmin.common.util.StringTool;
import com.khalil.kkadmin.module.ai.agent.chatsession.ChatSession;
import com.khalil.kkadmin.module.ai.agent.chatsession.ChatSessionMessage;
import com.khalil.kkadmin.module.ai.agent.chatsession.ChatSessionService;
import com.khalil.kkadmin.module.ai.agent.chatsession.JpaChatMemoryStore;
import com.khalil.kkadmin.module.ai.application.domain.AiApplication;
import com.khalil.kkadmin.module.ai.application.service.AiApplicationManager;
import com.khalil.kkadmin.module.ai.model.AiChatModelBuilder;
import com.khalil.kkadmin.module.ai.application.dto.ChatParams;
import com.khalil.kkadmin.module.ai.model.AiModel;
import com.khalil.kkadmin.module.ai.model.AiModelManager;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;

import static java.util.concurrent.TimeUnit.SECONDS;

@Slf4j
@Service
public class AiAgentServiceImpl implements AiAgentService {

    private final AiModelManager aiModelManager;
    private final AiApplicationManager aiApplicationManager;
    private final ChatSessionService chatSessionService;
    private final JpaChatMemoryStore jpaChatMemoryStore;

    public AiAgentServiceImpl(AiModelManager aiModelManager, AiApplicationManager aiApplicationManager, ChatSessionService chatSessionService, JpaChatMemoryStore jpaChatMemoryStore) {
        this.aiModelManager = aiModelManager;
        this.aiApplicationManager = aiApplicationManager;
        this.chatSessionService = chatSessionService;
        this.jpaChatMemoryStore = jpaChatMemoryStore;
    }

    @Override
    public void streamChat(ChatParams chatParams, Consumer<String> onData, Consumer<Throwable> onError, Consumer<Void> onComplete) {
        String username = chatParams.getUsername();
        Long sessionId = chatParams.getSessionId();
        Long applicationId = chatParams.getApplicationId();
        String chatParamsContent = chatParams.getContent();

        if (StringTool.isEmpty(chatParamsContent)) {
            log.error("聊天content参数为空！");
            onData.accept("问题为空！");
            onComplete.accept(null);
            return;
        }

        ChatSession chatSession;
        if (sessionId == null || sessionId < 0) {
            String sessionTitle = chatParamsContent.length() > 20 ? chatParamsContent.substring(0, 20) : chatParamsContent;
            chatSession = chatSessionService.createSession(username, applicationId, sessionTitle);
            sessionId = chatSession.getId();
            chatParams.setSessionId(sessionId);
        } else {
            chatSession = chatSessionService.getSession(sessionId);
            if (chatSession == null) {
                throw new RuntimeException("ChatSession not found");
            }
        }

        AiApplication aiApplication = aiApplicationManager.get(applicationId);
        AiModel aiModel = aiModelManager.get(aiApplication.getModelId());
        StreamingChatModel streamingChatModel = AiChatModelBuilder.buildStreamingChatModel(aiApplication, aiModel);
        if (streamingChatModel == null) {
            throw new RuntimeException("模型不存在");
        }

        // function call

        // 知识库搜索

        // rerank

        // 模型调用聊天
        ChatMemoryProvider chatMemoryProvider = (memoryId) -> MessageWindowChatMemory.builder().id(memoryId).maxMessages(10).chatMemoryStore(jpaChatMemoryStore).build();

        try {
            ChatAssistant chatAssistant = AiServices.builder(ChatAssistant.class).streamingChatModel(streamingChatModel).chatMemoryProvider(chatMemoryProvider).build();

            TokenStream tokenStream = chatAssistant.chat(chatParams.getSessionId(), chatParamsContent);

            CompletableFuture<ChatResponse> futureResponse = new CompletableFuture<>();

            ChatSessionMessage userSessionMessage = ChatSessionMessage.builder().role(ChatSessionMessage.ROLE_USER).content(chatParamsContent).build();
            userSessionMessage = chatSessionService.createSessionMessage(sessionId, userSessionMessage);

            final long finalSessionId = sessionId;
            final long finalUserSessionMessageId = userSessionMessage.getId();
            tokenStream.onPartialResponse(s -> {
                System.out.print(s);
                onData.accept(s);
            }).onCompleteResponse((ChatResponse chatResponse) -> {
                ChatSessionMessage assistantSessionMessage = ChatSessionMessage.builder().role(ChatSessionMessage.ROLE_ASSISTANT).content(chatResponse.aiMessage().text()).beforeId(finalUserSessionMessageId).build();
                chatSessionService.createSessionMessage(finalSessionId, assistantSessionMessage);
                futureResponse.complete(chatResponse);
            }).onError((Throwable throwable) -> {
                log.error("调用失败", throwable);
                ChatSessionMessage assistantSessionMessage = ChatSessionMessage.builder().role(ChatSessionMessage.ROLE_ASSISTANT).content("调用失败").beforeId(finalUserSessionMessageId).build();
                chatSessionService.createSessionMessage(finalSessionId, assistantSessionMessage);
                onError.accept(throwable);
                futureResponse.completeExceptionally(throwable);
            }).start();

            try {
                ChatResponse chatResponse = futureResponse.get(30, SECONDS);
                System.out.println("1111");
                System.out.println("\n" + chatResponse);
            } catch (InterruptedException | ExecutionException | TimeoutException e) {
                throw new RuntimeException(e);
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}
